import json
import os
from cryptography.fernet import Fernet
import base64
from pathlib import Path

class ConfigManager:
    def __init__(self):
        self.config_dir = os.path.join(os.path.expanduser("~"), ".vba_helper")
        self.config_file = os.path.join(self.config_dir, "config.enc")
        self.key_file = os.path.join(self.config_dir, ".key")
        self._ensure_config_dir()
        self._ensure_key()

    def _ensure_config_dir(self):
        """确保配置目录存在"""
        if not os.path.exists(self.config_dir):
            os.makedirs(self.config_dir)

    def _ensure_key(self):
        """确保加密密钥存在，如果不存在则创建"""
        if not os.path.exists(self.key_file):
            key = Fernet.generate_key()
            with open(self.key_file, "wb") as f:
                f.write(key)

    def _get_cipher(self):
        """获取加密器"""
        with open(self.key_file, "rb") as f:
            key = f.read()
        return Fernet(key)

    def save_api_keys(self, embedding_key: str, generation_key: str):
        """
        保存API密钥
        :param embedding_key: 阿里云嵌入服务的API密钥
        :param generation_key: DeepSeek生成服务的API密钥
        """
        config = {
            "embedding_api_key": embedding_key,
            "generation_api_key": generation_key
        }
        
        cipher = self._get_cipher()
        encrypted_data = cipher.encrypt(json.dumps(config).encode())
        
        with open(self.config_file, "wb") as f:
            f.write(encrypted_data)

    def load_api_keys(self):
        """
        加载API密钥
        :return: (embedding_key, generation_key) 元组，如果未设置则返回 (None, None)
        """
        if not os.path.exists(self.config_file):
            return None, None

        try:
            cipher = self._get_cipher()
            with open(self.config_file, "rb") as f:
                encrypted_data = f.read()
            
            decrypted_data = cipher.decrypt(encrypted_data)
            config = json.loads(decrypted_data.decode())
            
            return (
                config.get("embedding_api_key"),
                config.get("generation_api_key")
            )
        except Exception as e:
            print(f"加载配置时出错: {str(e)}")
            return None, None

    def clear_api_keys(self):
        """清除保存的API密钥"""
        if os.path.exists(self.config_file):
            os.remove(self.config_file) 